package com.kqzz.common.dynamicDB.config;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import javax.sql.DataSource;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.JndiDataSourceLookup;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

@SuppressWarnings("unchecked")
public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

	  @Nullable
	  private Map<Object, Object> targetDataSources;

	  @Nullable
	  private Object defaultTargetDataSource;

	  private boolean lenientFallback = true;

	  private DataSourceLookup dataSourceLookup = (DataSourceLookup)new JndiDataSourceLookup();

	  @Nullable
	  private Map<Object, DataSource> resolvedDataSources;

	  @Nullable
	  private DataSource resolvedDefaultDataSource;

	  public void setTargetDataSources(Map<Object, Object> targetDataSources) {
	    this.targetDataSources = targetDataSources;
	  }

	  public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
	    this.defaultTargetDataSource = defaultTargetDataSource;
	  }

	  public void setLenientFallback(boolean lenientFallback) {
	    this.lenientFallback = lenientFallback;
	  }

	  public void setDataSourceLookup(@Nullable DataSourceLookup dataSourceLookup) {
	    this.dataSourceLookup = (dataSourceLookup != null) ? dataSourceLookup : (DataSourceLookup)new JndiDataSourceLookup();
	  }

	  public void afterPropertiesSet() {
	    if (this.targetDataSources == null)
	      throw new IllegalArgumentException("Property 'targetDataSources' is required");
	    this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
	    this.targetDataSources.forEach((key, value) -> {
	          Object lookupKey = resolveSpecifiedLookupKey(key);
	          DataSource dataSource = resolveSpecifiedDataSource(value);
	          this.resolvedDataSources.put(lookupKey, dataSource);
	        });
	    if (this.defaultTargetDataSource != null)
	      this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
	  }

	  protected Object resolveSpecifiedLookupKey(Object lookupKey) {
	    return lookupKey;
	  }

	  protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
	    if (dataSource instanceof DataSource)
	      return (DataSource)dataSource;
	    if (dataSource instanceof String)
	      return this.dataSourceLookup.getDataSource((String)dataSource);
	    throw new IllegalArgumentException("Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
	  }

	  public Map<Object, DataSource> getResolvedDataSources() {
	    Assert.state((this.resolvedDataSources != null), "DataSources not resolved yet - call afterPropertiesSet");
	    return Collections.unmodifiableMap(this.resolvedDataSources);
	  }

	  @Nullable
	  public DataSource getResolvedDefaultDataSource() {
	    return this.resolvedDefaultDataSource;
	  }

	  public Connection getConnection() throws SQLException {
	    return determineTargetDataSource().getConnection();
	  }

	  public Connection getConnection(String username, String password) throws SQLException {
	    return determineTargetDataSource().getConnection(username, password);
	  }

	  public <T> T unwrap(Class<T> iface) throws SQLException {
	    if (iface.isInstance(this))
	      return (T)this;
	    return determineTargetDataSource().unwrap(iface);
	  }

	  public boolean isWrapperFor(Class<?> iface) throws SQLException {
	    return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
	  }

	  protected DataSource determineTargetDataSource() {
	    Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
	    Object lookupKey = determineCurrentLookupKey();
	    DataSource dataSource = this.resolvedDataSources.get(lookupKey);
	    if (dataSource == null && (this.lenientFallback || lookupKey == null))
	      dataSource = this.resolvedDefaultDataSource;
	    if (dataSource == null)
	      throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
	    return dataSource;
	  }

	  @Nullable
	  protected abstract Object determineCurrentLookupKey();
}
